import os
import tkinter as tk
from tkinter.filedialog import askopenfilename
from PIL import Image, ImageTk
import torch
from torchvision import transforms
from model import AlexNet
from networks.ClassicNetwork.ResNet import ResNet50


class MainWindow:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title('基于Resnet的野生动物识别')


        self.frame_left = tk.Frame(self.root, padx=1, pady=5, bg="#aaaaaa")
        self.frame_left.pack(padx=5, pady=10, fill="y", side=tk.TOP)

        # self.frame_left = tk.Frame(self.root, padx=1, pady=5, bg="#aaaaaa")
        # self.frame_left.pack(padx=5, pady=10, fill="y", side=tk.LEFT)
        photo = tk.PhotoImage(file="背景.png")
        self.frame_right = tk.Frame(self.root,image =  photo, padx=5, pady=5)
        self.frame_right.pack(padx=5, pady=10, fill="y", side=tk.BOTTOM)

        # self.button_loadModel = tk.Button(self.frame_right, text='请选择您想要识别的野生图片', command=self.loadModel, width=30, height=3)
        # self.button_loadModel.pack(fill="x")

        self.button_loadImage = tk.Button(self.frame_right, text='请选择您想要识别的野生图片',font=('宋体', 18), command=self.loadimg, width=30, height=2)
        self.button_loadImage.pack(fill="both")
        # self.button_loadImage.config(state=tk.DISABLED)

        self.button_predict = tk.Button(self.frame_right, text='该野生动物的类别属于',font=('宋体', 18), command=self.predict, width=30, height=2)
        self.button_predict.pack(fill="both")
        # self.button_predict.pack.grid(fill="both")
        # self.button_predict.config(state=tk.DISABLED)

        self.label_info = tk.Label(self.frame_right, font=('宋体', 12), justify=tk.LEFT, padx=2, pady=30)
        self.label_info.pack(fill="x")
        self.label_info.config(text="结果显示处")

        self.canvas = tk.Canvas(self.frame_left, bg='gray', height=300, width=400)
        self.canvas.pack(fill='x', expand='yes')

        self.root.mainloop()

    def load_image_to_canvas(self, file_path):
        """把给定路径的图像加载入self.img 并绘制到canvas"""
        """载入指定的模型"""
        try:
            default_dir = os.getcwd()
            modelPath = ""
            if modelPath == "":
                return
            self.label_info.config(text="")
            model = ResNet50(num_classes=7)
            model.load_state_dict(torch.load(modelPath))
            model.eval()
            self.model = model
        finally:
            self.button_loadImage.config(state=tk.NORMAL)
            self.label_info.config(text="")


        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)

        try:
            img = Image.open(file_path)
            self.img = img
            img_w, img_h = img.size
            if img_w > 400:
                img_w = 400
                img_h = img_h * (400 / img_w)
                img = resize(img_w, img_h, img)
            self.pil_img = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
            self.canvas.update()  # 获取宽高之前要先对于这个组件update()
            x, y = 0, (self.canvas.winfo_height() - img.size[1]) / 2
            self.canvas.create_image(x, y, anchor='nw', image=self.pil_img)
        except Exception as e:
            self.label_info.config(text="图片载入出错")
        finally:
            self.button_predict.config(state=tk.NORMAL)
            self.label_info.config(text="图片已载入\n点击预测按钮")

    def predict(self):
        """根据已载入的模型进行识别"""
        class_dict = {
            "0": "野生马",
            "1": "野生大象",
            "2": "蝴蝶",
            "3": "松鼠",
            "4": "野生牛",
            "5": "野生羊",
            "6": "蜘蛛",


        }

        data_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        preprocess_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        img = preprocess_transform(self.img)
        img.unsqueeze_(0)
        print(img.shape)
        with torch.no_grad():
            outputs, features = self.model(img)

            output = torch.squeeze(outputs)
            predict = torch.softmax(output, dim=0)
            predict_cla = torch.argmax(predict).numpy()  # 最大值位置索引
            print(predict_cla)
        class_str = class_dict[str(predict_cla)]
        prob_str = "%.1f" % (predict[predict_cla].item() * 100)
        self.label_info.config(text=f"类别：{class_str}\n可能性：{prob_str}%")

    # def loadModel(self):


    def loadimg(self):
        """载入指定的jpg图片"""
        default_dir = os.getcwd()
        photoPath = askopenfilename(title='打开一个照片（jpg格式）',
                                    initialdir=(os.path.expanduser(default_dir)),
                                    filetypes=[('jpg文件', '*.jpg'), ('All Files', '*')])
        if photoPath == "":
            return
        self.load_image_to_canvas(photoPath)


if __name__ == '__main__':
    win = MainWindow()